import input, preprocess, predict

def find_great_coach(predict_data, countries, coaches):
    # 遍历每个国家
    for country in countries:
        # 遍历每个教练
        for coach in coaches:
            # 如果教练的国家与当前国家相同
            if coach['Country'] == country:
                # 根据教练的实力增加总奖牌数和金牌数
                predict_data.loc[predict_data['Country'] == country, 'Total_Medals'] += coach['Strength']
                predict_data.loc[predict_data['Country'] == country, 'Gold_Medals'] += coach['Strength'] * 0.5  # 假设金牌数增加为实力的一半

    return predict_data

# 主程序
if __name__ == "__main__":
    file_path = "./data/"  # 替换为你的文件路径
    dataframes = input.load_data(file_path)
    # 特征提取
    features = preprocess.preprocess_data(dataframes)

    # 输入需要预测的年份，和主办国家
    year_to_predict = 2028; host_country="United States"; num_of_programs = 350

    # 预测
    predict_data = predict.train_and_predict(features, year_to_predict, host_country, num_of_programs)

    # 需要寻找great_coach的国家
    countries = ["United States", "China", "France"]

    # great_coach的基本属性，参照Lang Ping和Bela Karolyi的数据
    # Sport为具体运动项目； Type为教练所教的为集体项目（C），个人项目（P），既有集体又有个人（B）；Strength为教练实力
    coaches = [
        {"Coach_Name": "Coach_A", "Sport": "Table Tennis", "Type": "C", "Strength": 0.7, "Country": "United States"},
        {"Coach_Name": "Coach_B", "Sport": "Basketball", "Type": "P", "Strength": 4.5, "Country": "China"},
        {"Coach_Name": "Coach_C", "Sport": "Volleyball", "Type": "B", "Strength": 2.2, "Country": "France"}
    ]

    # 在2028年寻找教练
    predictions = find_great_coach(predict_data, countries, coaches)

    # 评价
    predictions = predict.evaluate_strength(predict_data, features, year_to_predict)

    # 过滤其他预测结果，只保留指定国家的数据
    filtered_predictions = predictions[predictions['Country'].isin(countries)]

    # 输出结果
    print("\nPredicted Medals for", year_to_predict, "Olympics:")
    # print(predictions.to_string(index=False))
    print(filtered_predictions)